package algorithm

import (
	"bytes"
	"crypto/rand"
	"crypto/rsa"
	"encoding/base64"
)

var (
	privateKey *rsa.PrivateKey
	publicKey  *rsa.PublicKey
)

// 初始化 RSA 密钥对
func RSAInit() {
	var err error
	privateKey, err = rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		panic(err)
	}
	publicKey = &privateKey.PublicKey
}

// RSA加密：对长字符串进行分段加密
func RSAEncryptLongMessage(utf8Str string) ([]string, error) {
	plainData := []byte(utf8Str)
	// RSA 每次能加密的最大数据长度: keySize - 11（PKCS#1 v1.5 padding）
	chunkSize := publicKey.Size() - 11
	var encryptedChunks []string

	for i := 0; i < len(plainData); i += chunkSize {
		end := i + chunkSize
		if end > len(plainData) {
			end = len(plainData)
		}
		chunk := plainData[i:end]
		encryptedChunk, err := rsa.EncryptPKCS1v15(rand.Reader, publicKey, chunk)
		if err != nil {
			return nil, err
		}
		encryptedChunks = append(encryptedChunks, base64.StdEncoding.EncodeToString(encryptedChunk))
	}
	return encryptedChunks, nil
}

// RSA解密：将密文块逐个解密拼接成原始数据
func RSADecryptLongMessage(encryptedChunks []string) (string, error) {
	var buffer bytes.Buffer
	for _, encodedChunk := range encryptedChunks {
		cipherChunk, err := base64.StdEncoding.DecodeString(encodedChunk)
		if err != nil {
			return "", err
		}
		decryptedChunk, err := rsa.DecryptPKCS1v15(rand.Reader, privateKey, cipherChunk)
		if err != nil {
			return "", err
		}
		buffer.Write(decryptedChunk)
	}
	return string(buffer.Bytes()), nil
}
